"""
PCMCIPlusEngine 模块 - 基于PCMCI+算法的因果发现

PCMCI+ 是PCMCI的改进版本，能够识别完整的滞后和同期因果图。相比标准的PCMCI，
PCMCIplus能够在因果充分性、忠实性和马尔可夫条件的标准假设下，识别同期因果关系
(直到同期链接的马尔可夫等价类)。
"""

import numpy as np
import pandas as pd
import logging
from typing import Dict, List, Tuple
import tigramite.data_processing as pp
from tigramite.pcmci import PCMCI

from .PCMCIEngine import PCMCIEngine

logger = logging.getLogger(__name__)


class PCMCIPlusEngine(PCMCIEngine):
    """基于PCMCI+算法的因果分析引擎"""

    def __init__(self, config):
        """
        初始化PCMCI+因果分析引擎
        Args:
            config: 配置对象，继承PCMCIEngine的所有配置
        """
        super().__init__(config)

    def _run_algorithm(self, data: pd.DataFrame) -> Dict:
        """
        执行PCMCI+算法
        Args:
            data: 输入数据
        Returns:
            包含检验结果的字典
        """
        try:
            # 创建Tigramite数据框架
            dataframe = pp.DataFrame(
                data.values, var_names=data.columns, datatime=data.index.values
            )

            # 创建PCMCI对象并执行分析
            pcmci = PCMCI(
                dataframe=dataframe, cond_ind_test=self.cond_ind_test, verbosity=0
            )

            # 执行PCMCI+分析
            results = pcmci.run_pcmciplus(
                tau_min=self.tau_min, tau_max=self.tau_max, pc_alpha=self.alpha
            )

            return {
                "val_matrix": results["val_matrix"],
                "p_matrix": results["p_matrix"],
                "alpha": self.alpha,
            }
        except Exception as e:
            logger.error(f"PCMCI+分析失败: {str(e)}")
            n_vars = data.shape[1]
            return {
                "val_matrix": np.zeros((n_vars, n_vars, self.tau_max)),
                "p_matrix": np.ones((n_vars, n_vars, self.tau_max)),
                "alpha": self.alpha,
            }

    def _get_method_name(self) -> str:
        """获取算法方法名称"""
        return "PCMCIplus"
